structure StratCore =
struct
val _ = Compiler.Control.Print.printDepth := 100;

(*====================================================================================================================*)
datatype ('i,'o) options  = None of 'i | Some of 'o;
type     ('i,'o) strategy = 'i -> ('i,'o) options;

(*====================================================================================================================*)
datatype term      = T of string * term list
                   | V of string;
datatype stratcore = Id
                   | Fail
                   | Rule   of term  * term
                   | Seq    of stratcore * stratcore
                   | Choice of stratcore * stratcore;

(*====================================================================================================================*)
(* Standard rewriting from "Term Rewriting and All That" (p.81, Fig.4.7).           *)
(* Note: some functions are generalized to be polymorphic for both terms and types. *)
structure R:
  sig
    val rewrite: term * term -> term -> term
    val zip: 'a list * 'b list -> ('a * 'b) list
    val indom: string -> (string * 'a) list -> bool
    val app: (string * 'a) list -> string -> 'a
  end
= struct
type 'a subst = (string * 'a) list;

val indom: string -> 'a subst -> bool
  = fn x => fn s => List.exists (fn (y,_) => x = y) s;

val rec app: 'a subst -> string -> 'a
  = fn ((y,t)::s) => (fn x => if x = y then t else app s x)
     | []         =>  fn x => raise Match;

val rec zip: 'a list * 'b list -> ('a * 'b) list
  = fn ([],    []   ) => []
     | (x::xs, y::ys) => (x,y)::zip(xs,ys)
     | _              => raise Match

val rec lift: term subst -> term -> term
  = fn s => (fn V x      => if indom x s then app s x else V x
              | T (f,ts) => T(f, map (lift s) ts)) 

val rec matchs: (term * term) list * term subst -> term subst
  = fn ([]                   , s) => s
     | ((V x, t)::S          , s) => if indom x s then if app s x = t then matchs(S, s) else raise Match
                                     else matchs(S, (x,t)::s)
     | ((t, V _)::S          , s) => raise Match
     | ((T(f,ts), T(g,us))::S, s) => if f = g then matchs(zip(ts,us)@S, s) else raise Match

val match: term * term -> term subst
  = fn (pat, obj) => matchs ([(pat, obj)], []);

val rewrite: term * term -> term -> term
  = fn (l,r) => fn t => lift (match(l,t)) r

end;

(*====================================================================================================================*)
(* Semantics of StratCore *)
val rec interpret: stratcore -> (term, term) strategy
  = fn Id             => (fn t => Some t)
     | Fail           => (fn t => None t)
     | Rule (l, r)    => (fn t => Some (R.rewrite (l, r) t) handle _ => None t)
     | Seq (s, s')    => (fn Some t => interpret s' t
                           | None t => None t)
                          o
                         (interpret s)
     | Choice (s, s') => (fn t => case (interpret s t) of
                                    Some t' => Some t'
                                  | None t' => interpret s' t)



(*====================================================================================================================*)
(*====================================================================================================================*)
(*====================================================================================================================*)
(* Type analysis of StratCore *)
datatype tyT = TyTerm of string * tyT list | TyVar  of int;
datatype tyS = TyRule of tyT * tyT;
type context = (string * tyT) list;

val nextSym = ref 0
val nextTyVar: unit -> tyT
  = fn () => let val i = !nextSym
             in  nextSym := i + 1; TyVar i
             end

(*====================================================================================================================*)
(* Standard unification from "Term Rewriting and All That" (p.80, Fig.4.5) *)
(* Note: the only difference here is that terms are replaced by types.     *)
structure U:
  sig
    val unify: tyT * tyT -> context
    val lift: context -> tyT -> tyT
    val solve: (tyT * tyT) list * context -> context
  end
= struct

val rec lift: context -> tyT -> tyT
  = fn s => (fn TyVar  x      => if R.indom (Int.toString x) s then R.app s (Int.toString x) else TyVar x
              | TyTerm (f,ts) => TyTerm (f, map (lift s) ts))

val rec occurs: string -> tyT -> bool
  = fn x => (fn TyVar  y      => x = (Int.toString y)
              | TyTerm (_,ts) => List.exists (occurs x) ts)

val rec solve: (tyT * tyT) list * context -> context
  = fn ([]               , s) => s
     | ((TyVar x, t)  ::S, s) => if TyVar x = t then solve (S,s) else elim(Int.toString x,t,S,s)
     | ((t, TyVar x)  ::S, s) => elim(Int.toString x,t,S,s)
     | ((TyTerm(f,ts),
         TyTerm(g,us))::S, s) => if f = g then solve(R.zip(ts,us)@S, s) else raise Match

and elim: string * tyT * (tyT * tyT) list * context -> context
  = fn (x,t,S,s) => if occurs x t then raise Match
                    else let val xt = lift [(x,t)]
                         in  solve(map (fn (t1,t2) => (xt t1, xt t2)) S,
                                   (x,t)::(map (fn (y,u) => (y, xt u)) s))
                         end;

val unify: tyT * tyT -> context
  = fn (t1,t2) => solve([(t1,t2)], []) 
end;



(*====================================================================================================================*)
exception FreeVarError;
exception DeadCodeError;
exception FailingStrategyError;
exception ApplicationError;

val rec subtype: tyT * tyT -> bool
  = fn (_            , TyVar _      ) => true
     | (TyTerm (x,xs), TyTerm (y,ys)) => x = y andalso subtypeL (xs, ys)
     | (x            , y            ) => if x = y then true else false

and subtypeL: tyT list * tyT list -> bool
  = fn ([]   , []   ) => true
     | ([]   , _    ) => false
     | (_    , []   ) => false
     | (x::xs, y::ys) => subtype (x, y) andalso subtypeL (xs, ys)

val rec isReachable: tyS list * tyS list -> bool
  = fn ([]              , _               ) => true
     | (_               , []              ) => true
     | (TyRule (a,b)::xs, TyRule (c,d)::ys) =>
         not (subtype (c,a)) andalso
         isReachable (TyRule (a,b)::[], ys) andalso
         isReachable (xs, TyRule (c,d)::ys)

val rec getStrategies: tyS list * tyS list -> tyS list
  = fn ([]              , _               ) => []
     | (_               , []              ) => []
     | (TyRule(a, b)::xs, TyRule(c, d)::ys) =>
         let val ctx' = (SOME (U.unify (b, c))) handle _ => NONE;
         in  (case ctx' of
                SOME ctx => [TyRule(U.lift ctx a, U.lift ctx d)]
              | NONE     => []) @
             getStrategies (TyRule (a, b)::[], ys) @
             getStrategies (xs, TyRule (c, d)::ys)
         end


val rec update: context -> term -> context
  = fn ctx => (fn V x     => if R.indom x ctx then ctx
                             else (x, nextTyVar())::ctx
                | T(_,ts) => foldl (fn (t,ctx') => update ctx' t) ctx ts);

val rec query: string -> context -> tyT
  = fn x => (fn ((y,t)::ctx) => if x = y then t else query x ctx
              | []           => raise FreeVarError)

val rec typeOfT: context -> term -> tyT = fn ctx =>
 fn V x           => query x ctx
  | T (f,ts)      => TyTerm(f, map (typeOfT ctx) ts)

and typeOfS: context -> stratcore -> tyS list = fn ctx =>
 fn Id            => let val t = nextTyVar()
                     in  [TyRule (t, t)]
                     end
  | Fail          => []
  | Rule   (l, r) => let val ctx' = update ctx l
                         val tL = typeOfT ctx' l
                         val tR = typeOfT ctx' r
                     in  [TyRule (tL, tR)]
                     end
  | Choice (l, r) => let val sL = typeOfS ctx l
                         val sR = typeOfS ctx r
                     in  case isReachable (sL, sR) of
                           true  => sL @ sR
                         | false => raise DeadCodeError
                     end
  | Seq    (l, r) => let val sL = typeOfS ctx l
                         val sR = typeOfS ctx r
                     in  case getStrategies (sL, sR) of
                           [] => raise FailingStrategyError
                         | ss => ss
                     end


val rec reduce: tyT -> tyS list -> tyT list = fn tt =>
  fn []               => []
   | TyRule (l,r)::ss => let val ctx' = SOME (U.unify (l, tt)) handle _ => NONE
                         in  (case ctx' of
                                SOME ctx => [U.lift ctx r]
                              | NONE     => [])
                             @
                             (reduce tt ss) 
                         end

val rec typeCheck: stratcore -> term -> tyT list
  = fn s => fn t => let val ss = typeOfS [] s;
                        val tt = typeOfT [] t;
                    in  case (reduce tt ss) of
                          []  => raise ApplicationError
                        | tts => tts
                    end

(*====================================================================================================================*)
(* tests *)
val timesToPlus
  = Rule(T("Plus", [V "e",
                    T("Mult", [V "t",
                               T("Num", [V "x"]) ]) ]),
         T("Plus", [T("Plus", [V "e", V "t"]),
                    T("Mult", [V "t",
                               T("Num", [V "x",
                                         T("-", []),
                                         T("1", []) ]) ]) ]) );
val timesToPlus'
  = Rule(T("Plus", [V "e",
                    T("Mult", [V "t",
                              T("Num", [T("2", [])]) ]) ]),
         T("Plus", [T("Plus", [V "e", V "t"]),
                    V "t"]) );
val timesToPluses = Choice (timesToPlus, timesToPlus');
(*val testChoice = typeOfS [] timesToPluses; (* this raises a DeadCodeError *)*)

val timesZero
  = Rule(T("Mult", [T("Term", [T("Num", [T("0", []) ]) ]),
                    V "x"]),
         T("Term", [T("Num", [T("0", []) ]) ]) );
val timesZero'
  = Rule(T("Mult", [V "x",
                    T("Num", [T("0", []) ]) ]),
         T("Term", [T("Num", [T("0", []) ]) ]) );
val timesZeroes = Choice (timesZero, timesZero');

val timesOne
  = Rule(T("Mult", [T("Term", [T("Num", [T("1", []) ]) ]),
                    V "x"]),
         T("Term", [V "x"]) );
val timesOne'
  = Rule(T("Mult", [V "x",
                    T("Num", [T("1", []) ]) ]),
         V "x" );
val timesOnes = Choice (timesOne, timesOne');
val simpTimes = Seq (timesZeroes, timesOnes);
(*val testSeq   = typeOfS [] simpTimes; (* this raises FailingStrategyError *)*)

val r1 = Rule(T("Expr", [V "x",
                         T("+", []),
                         T("0", []) ]),
              V "x" );
val testApp1 = typeCheck r1 (T("Expr", [T("Expr", [T("a", []) ]),
                                        T("+", []),
                                        T("0", []) ]) );
    (*=[TyTerm ("Expr", [TyTerm ("a", []) ])]*)
(*val testApp2
  = typeCheck r1 (T("Expr", [T("Expr", [T("a", []) ]),
                             T("+", []),
                             T("1", []) ]) ); (* this raises ApplicationError *)*)

(*====================================================================================================================*)
end;
